#!/bin/bash

MODEL="$1"
REDIS_KEY="$2"

if [ -z "$MODEL" ] || [ -z "$REDIS_KEY" ]; then
    echo "Usage: $0 <model> <redis_key>" >&2
    exit 1
fi

PROMPT=$(cat)
if [ -z "$PROMPT" ]; then
    echo "No prompt provided via stdin" >&2
    exit 1
fi

# API 配置
if [ "$MODEL" == "deepseek" ]; then
    API_URL="https://api.deepseek.com/v1/chat/completions"
    API_KEY="sk-4222ff3eeb5f498c9a703dc83df73975"
    MODEL_NAME="deepseek-chat"
elif [ "$MODEL" == "bailian" ]; then
    API_URL="https://api.bailian.ai/v1/chat/completions"
    API_KEY="$BAILIAN_API_KEY"
    MODEL_NAME="bailian-chat"
else
    echo "Unsupported model: $MODEL" >&2
    exit 1
fi

[ -z "$API_KEY" ] && { echo "API key not set for $MODEL" >&2; exit 1; }

MAX_TOKENS=4000
TEMPERATURE=0.7

JSON_DATA=$(jq -n \
  --arg model "$MODEL_NAME" \
  --arg prompt "$PROMPT" \
  --argjson stream true \
  --argjson max_tokens $MAX_TOKENS \
  --argjson temperature $TEMPERATURE \
  '{
    model: $model,
    messages: [{role: "user", content: $prompt}],
    stream: $stream,
    max_tokens: $max_tokens,
    temperature: $temperature
  }')

curl -s -N -X POST "$API_URL" \
  -H "Authorization: Bearer $API_KEY" \
  -H "Content-Type: application/json" \
  --data-binary "$JSON_DATA" \
| while IFS= read -r line; do
    [[ -z "$line" ]] && continue
    [[ "$line" != data:* ]] && continue
    data="${line#data: }"
    [[ "$data" == "[DONE]" ]] && break
    content=$(echo "$data" | jq -r '.choices[0].delta.content // ""' 2>/dev/null)
    [[ -n "$content" && "$content" != "null" ]] && redis-cli XADD "$REDIS_KEY" "*" "message" "$content" > /dev/null 2>&1
done

# 写入结束标记
redis-cli XADD "$REDIS_KEY" "*" "message" "__END__" > /dev/null 2>&1
